from wavegan import WaveGANGenerator
import argparse
import os
import numpy as np
import torchvision.transforms as transforms
from torchvision.utils import save_image
from torch.utils.data import DataLoader
from torchvision import datasets
from torch.autograd import Variable
import torch.nn as nn
import torch
from torch.utils import data
from PIL import Image
from utils import *
import datetime

args = parse_arguments()
cuda = True if torch.cuda.is_available() else False

def eval():
    output_dir = args['output_dir']
    netG_path = os.path.join(output_dir, "generator_500`.pkl")
    # Load the generator model.
    if os.path.exists(netG_path):
        print("Loading generator model from: {}".format(netG_path))
        generator = WaveGANGenerator()
        generator.load_state_dict(torch.load(netG_path))
    else:
        print("Generator model not found.")
        return  
    # Sample noise used for generated output.
    sample_noise = torch.randn(args['sample_size'], args['latent_dim'])
    sample_noise_Var = autograd.Variable(sample_noise, requires_grad=False)
    LOGGER.info("Generating samples...")
    sample_out = generator(sample_noise_Var)
    if cuda:
        sample_out = sample_out.cpu()
    sample_out = sample_out.data.numpy()
    save_samples(sample_out, "test", output_dir)
    
    


if __name__ == '__main__':
    eval()
    